# borrowed from "https://github.com/marvis/pytorch-mobilenet"

import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict


class VGG16(nn.Module):
    def __init__(self, cfg=None, batch_norm=True):
        super(VGG16, self).__init__()

        layers = []
        in_channels = 3
        # for v in cfg:
        #     if v == 'M':
        #         layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        #     elif v == 'C':
        #         layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
        #     else:
        #         conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
        #         if batch_norm:
        #             layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
        #         else:
        #             layers += [conv2d, nn.ReLU(inplace=True)]
        #         in_channels = v
        model_dict = {}
        for i, v in enumerate(cfg):
            if v == 'M':
                model_dict[str(i)] = nn.MaxPool2d(kernel_size=2, stride=2)
            elif v == 'C':
                model_dict[str(i)] = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
            else:
                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
                if batch_norm:
                    model_dict["conv2d_" + str(i)] = conv2d
                    model_dict["bn_" + str(i)] = nn.BatchNorm2d(v)
                    model_dict["relu_" + str(i)] = nn.ReLU(inplace=True)
                else:
                    model_dict["conv2d_" + str(i)] = conv2d
                    model_dict["relu_" + str(i)] = nn.ReLU(inplace=True)
                in_channels = v

        model_dict["pool5"] = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        model_dict["conv6"] = nn.Conv2d(512, 1024, kernel_size=3, padding=1)
        # conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6)
        # conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=1)
        model_dict["relu_after6"] = nn.ReLU(inplace=True)
        model_dict["conv7"] = nn.Conv2d(1024, 1024, kernel_size=1)
        model_dict["relu_after7"] = nn.ReLU(inplace=True)

        self.model = nn.Sequential(
            OrderedDict(model_dict)
        )

    def forward(self, x):
        x = self.model(x)
        return x
